"""
Health monitoring and performance tracking for CodeMCP servers.

Provides health checks, performance metrics, and system monitoring
capabilities for all MCP servers in the CodeMCP framework.
"""

import asyncio
import time
import psutil
import json
from typing import Dict, Any, List, Optional, Callable
from datetime import datetime, timedelta
from dataclasses import dataclass, field
from enum import Enum
from .error_handler import log_info, log_warning, log_debug


class HealthStatus(Enum):
    """Health status levels."""
    HEALTHY = "healthy"
    WARNING = "warning" 
    CRITICAL = "critical"
    UNKNOWN = "unknown"


@dataclass
class MetricPoint:
    """A single metric measurement."""
    timestamp: datetime
    value: float
    tags: Dict[str, str] = field(default_factory=dict)
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "timestamp": self.timestamp.isoformat(),
            "value": self.value,
            "tags": self.tags
        }


@dataclass
class HealthCheck:
    """Definition of a health check."""
    name: str
    function: Callable
    interval_seconds: int = 60
    timeout_seconds: int = 30
    critical_threshold: Optional[float] = None
    warning_threshold: Optional[float] = None
    enabled: bool = True
    
    def __post_init__(self):
        self.last_run: Optional[datetime] = None
        self.last_result: Optional[Dict[str, Any]] = None
        self.consecutive_failures: int = 0


class MetricsCollector:
    """Collects and stores performance metrics."""
    
    def __init__(self, max_points: int = 1000):
        self.max_points = max_points
        self.metrics: Dict[str, List[MetricPoint]] = {}
        self.start_time = datetime.utcnow()
    
    def record_metric(self, 
                     name: str, 
                     value: float, 
                     tags: Dict[str, str] = None):
        """Record a metric value."""
        if name not in self.metrics:
            self.metrics[name] = []
        
        point = MetricPoint(
            timestamp=datetime.utcnow(),
            value=value,
            tags=tags or {}
        )
        
        self.metrics[name].append(point)
        
        # Keep only the most recent points
        if len(self.metrics[name]) > self.max_points:
            self.metrics[name] = self.metrics[name][-self.max_points:]
    
    def get_metric_history(self, 
                          name: str, 
                          duration_minutes: int = 60) -> List[MetricPoint]:
        """Get metric history for the specified duration."""
        if name not in self.metrics:
            return []
        
        cutoff_time = datetime.utcnow() - timedelta(minutes=duration_minutes)
        return [point for point in self.metrics[name] 
                if point.timestamp >= cutoff_time]
    
    def get_latest_metric(self, name: str) -> Optional[MetricPoint]:
        """Get the most recent value for a metric."""
        if name not in self.metrics or not self.metrics[name]:
            return None
        return self.metrics[name][-1]
    
    def get_metric_summary(self, 
                          name: str, 
                          duration_minutes: int = 60) -> Dict[str, Any]:
        """Get statistical summary of a metric."""
        history = self.get_metric_history(name, duration_minutes)
        
        if not history:
            return {"name": name, "count": 0}
        
        values = [point.value for point in history]
        
        return {
            "name": name,
            "count": len(values),
            "min": min(values),
            "max": max(values),
            "mean": sum(values) / len(values),
            "latest": values[-1],
            "duration_minutes": duration_minutes
        }
    
    def get_all_metrics(self) -> Dict[str, Any]:
        """Get all metric summaries."""
        return {
            name: self.get_metric_summary(name)
            for name in self.metrics.keys()
        }


class HealthMonitor:
    """Health monitoring system for MCP servers."""
    
    def __init__(self, server_name: str):
        self.server_name = server_name
        self.metrics = MetricsCollector()
        self.health_checks: Dict[str, HealthCheck] = {}
        self.current_status = HealthStatus.UNKNOWN
        self.is_running = False
        self.monitor_task: Optional[asyncio.Task] = None
        
        # Register default health checks
        self._register_default_health_checks()
    
    def _register_default_health_checks(self):
        """Register default health checks."""
        # System resource checks
        self.register_health_check(
            "memory_usage",
            self._check_memory_usage,
            interval_seconds=30,
            warning_threshold=80.0,  # 80% memory usage
            critical_threshold=95.0  # 95% memory usage
        )
        
        self.register_health_check(
            "cpu_usage",
            self._check_cpu_usage,
            interval_seconds=30,
            warning_threshold=80.0,  # 80% CPU usage
            critical_threshold=95.0  # 95% CPU usage
        )
        
        self.register_health_check(
            "disk_usage",
            self._check_disk_usage,
            interval_seconds=60,
            warning_threshold=80.0,  # 80% disk usage
            critical_threshold=90.0  # 90% disk usage
        )
        
        # Application-specific checks
        self.register_health_check(
            "response_time",
            self._check_response_time,
            interval_seconds=60,
            warning_threshold=2.0,   # 2 second response time
            critical_threshold=5.0   # 5 second response time
        )
    
    def register_health_check(self,
                             name: str,
                             function: Callable,
                             interval_seconds: int = 60,
                             timeout_seconds: int = 30,
                             warning_threshold: Optional[float] = None,
                             critical_threshold: Optional[float] = None,
                             enabled: bool = True):
        """Register a new health check."""
        health_check = HealthCheck(
            name=name,
            function=function,
            interval_seconds=interval_seconds,
            timeout_seconds=timeout_seconds,
            warning_threshold=warning_threshold,
            critical_threshold=critical_threshold,
            enabled=enabled
        )
        
        self.health_checks[name] = health_check
        log_info(f"Registered health check: {name} for {self.server_name}")
    
    def unregister_health_check(self, name: str):
        """Unregister a health check."""
        if name in self.health_checks:
            del self.health_checks[name]
            log_info(f"Unregistered health check: {name} for {self.server_name}")
    
    async def start(self):
        """Start the health monitoring system."""
        if self.is_running:
            return
        
        self.is_running = True
        self.monitor_task = asyncio.create_task(self._monitoring_loop())
        log_info(f"Started health monitoring for {self.server_name}")
    
    async def stop(self):
        """Stop the health monitoring system."""
        if not self.is_running:
            return
        
        self.is_running = False
        
        if self.monitor_task:
            self.monitor_task.cancel()
            try:
                await self.monitor_task
            except asyncio.CancelledError:
                pass
        
        log_info(f"Stopped health monitoring for {self.server_name}")
    
    async def _monitoring_loop(self):
        """Main monitoring loop."""
        while self.is_running:
            try:
                await self._run_health_checks()
                await self._collect_metrics()
                await self._update_overall_status()
                
                # Sleep for a short interval
                await asyncio.sleep(10)
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                log_warning(f"Error in monitoring loop for {self.server_name}: {e}")
                await asyncio.sleep(30)  # Wait longer on error
    
    async def _run_health_checks(self):
        """Run all enabled health checks."""
        current_time = datetime.utcnow()
        
        for check_name, health_check in self.health_checks.items():
            if not health_check.enabled:
                continue
            
            # Check if it's time to run this health check
            if (health_check.last_run and 
                (current_time - health_check.last_run).total_seconds() < health_check.interval_seconds):
                continue
            
            try:
                # Run the health check with timeout
                result = await asyncio.wait_for(
                    self._run_single_health_check(health_check),
                    timeout=health_check.timeout_seconds
                )
                
                health_check.last_result = result
                health_check.consecutive_failures = 0
                
            except asyncio.TimeoutError:
                log_warning(f"Health check '{check_name}' timed out for {self.server_name}")
                health_check.consecutive_failures += 1
                health_check.last_result = {
                    "status": HealthStatus.CRITICAL.value,
                    "message": "Health check timed out",
                    "value": None
                }
                
            except Exception as e:
                log_warning(f"Health check '{check_name}' failed for {self.server_name}: {e}")
                health_check.consecutive_failures += 1
                health_check.last_result = {
                    "status": HealthStatus.CRITICAL.value,
                    "message": str(e),
                    "value": None
                }
            
            health_check.last_run = current_time
    
    async def _run_single_health_check(self, health_check: HealthCheck) -> Dict[str, Any]:
        """Run a single health check."""
        if asyncio.iscoroutinefunction(health_check.function):
            result = await health_check.function()
        else:
            result = health_check.function()
        
        # Determine status based on thresholds
        status = HealthStatus.HEALTHY
        if isinstance(result, (int, float)):
            value = result
            message = f"Value: {value}"
            
            if (health_check.critical_threshold is not None and 
                value >= health_check.critical_threshold):
                status = HealthStatus.CRITICAL
            elif (health_check.warning_threshold is not None and 
                  value >= health_check.warning_threshold):
                status = HealthStatus.WARNING
                
        elif isinstance(result, dict):
            value = result.get("value")
            message = result.get("message", "")
            status_str = result.get("status", HealthStatus.HEALTHY.value)
            status = HealthStatus(status_str)
        else:
            value = None
            message = str(result)
        
        return {
            "status": status.value,
            "value": value,
            "message": message,
            "timestamp": datetime.utcnow().isoformat()
        }
    
    async def _collect_metrics(self):
        """Collect system and application metrics."""
        current_time = datetime.utcnow()
        
        # System metrics
        try:
            # Memory usage
            memory = psutil.virtual_memory()
            self.metrics.record_metric("memory_usage_percent", memory.percent)
            self.metrics.record_metric("memory_used_mb", memory.used / 1024 / 1024)
            
            # CPU usage
            cpu_percent = psutil.cpu_percent(interval=None)
            self.metrics.record_metric("cpu_usage_percent", cpu_percent)
            
            # Disk usage
            disk = psutil.disk_usage('/')
            disk_percent = (disk.used / disk.total) * 100
            self.metrics.record_metric("disk_usage_percent", disk_percent)
            
            # Process-specific metrics
            process = psutil.Process()
            self.metrics.record_metric("process_memory_mb", 
                                     process.memory_info().rss / 1024 / 1024)
            self.metrics.record_metric("process_cpu_percent", 
                                     process.cpu_percent())
            
        except Exception as e:
            log_warning(f"Failed to collect system metrics for {self.server_name}: {e}")
    
    async def _update_overall_status(self):
        """Update the overall health status."""
        if not self.health_checks:
            self.current_status = HealthStatus.UNKNOWN
            return
        
        # Count statuses
        status_counts = {status: 0 for status in HealthStatus}
        
        for health_check in self.health_checks.values():
            if not health_check.enabled or not health_check.last_result:
                continue
            
            status_str = health_check.last_result.get("status", HealthStatus.UNKNOWN.value)
            status = HealthStatus(status_str)
            status_counts[status] += 1
        
        # Determine overall status
        if status_counts[HealthStatus.CRITICAL] > 0:
            self.current_status = HealthStatus.CRITICAL
        elif status_counts[HealthStatus.WARNING] > 0:
            self.current_status = HealthStatus.WARNING
        elif status_counts[HealthStatus.HEALTHY] > 0:
            self.current_status = HealthStatus.HEALTHY
        else:
            self.current_status = HealthStatus.UNKNOWN
    
    # Default health check implementations
    
    def _check_memory_usage(self) -> float:
        """Check system memory usage."""
        return psutil.virtual_memory().percent
    
    def _check_cpu_usage(self) -> float:
        """Check system CPU usage."""
        return psutil.cpu_percent(interval=1)
    
    def _check_disk_usage(self) -> float:
        """Check disk usage."""
        disk = psutil.disk_usage('/')
        return (disk.used / disk.total) * 100
    
    def _check_response_time(self) -> Dict[str, Any]:
        """Check average response time."""
        # This would be implemented by the specific MCP server
        # For now, return a placeholder
        return {
            "status": HealthStatus.HEALTHY.value,
            "value": 0.5,  # 500ms
            "message": "Average response time: 500ms"
        }
    
    # Public interface methods
    
    def get_current_health(self) -> Dict[str, Any]:
        """Get current health status."""
        return {
            "server_name": self.server_name,
            "overall_status": self.current_status.value,
            "timestamp": datetime.utcnow().isoformat(),
            "uptime_seconds": (datetime.utcnow() - self.metrics.start_time).total_seconds()
        }
    
    async def get_health_status(self) -> Dict[str, Any]:
        """Get comprehensive health status."""
        health_checks_status = {}
        
        for name, health_check in self.health_checks.items():
            health_checks_status[name] = {
                "enabled": health_check.enabled,
                "last_run": health_check.last_run.isoformat() if health_check.last_run else None,
                "last_result": health_check.last_result,
                "consecutive_failures": health_check.consecutive_failures,
                "interval_seconds": health_check.interval_seconds
            }
        
        return {
            "server_name": self.server_name,
            "overall_status": self.current_status.value,
            "health_checks": health_checks_status,
            "metrics_summary": self.metrics.get_all_metrics(),
            "timestamp": datetime.utcnow().isoformat(),
            "monitoring_active": self.is_running
        }
    
    def get_metrics_history(self, 
                           metric_name: str, 
                           duration_minutes: int = 60) -> List[Dict[str, Any]]:
        """Get historical data for a specific metric."""
        history = self.metrics.get_metric_history(metric_name, duration_minutes)
        return [point.to_dict() for point in history]
    
    def record_custom_metric(self, 
                            name: str, 
                            value: float, 
                            tags: Dict[str, str] = None):
        """Record a custom metric value."""
        self.metrics.record_metric(name, value, tags)
        log_debug(f"Recorded custom metric {name}={value} for {self.server_name}")
    
    def get_health_summary(self) -> Dict[str, Any]:
        """Get a summary of health status for dashboards."""
        enabled_checks = sum(1 for check in self.health_checks.values() if check.enabled)
        healthy_checks = sum(1 for check in self.health_checks.values() 
                           if check.enabled and check.last_result and 
                           check.last_result.get("status") == HealthStatus.HEALTHY.value)
        
        latest_memory = self.metrics.get_latest_metric("memory_usage_percent")
        latest_cpu = self.metrics.get_latest_metric("cpu_usage_percent")
        
        return {
            "server_name": self.server_name,
            "status": self.current_status.value,
            "health_score": (healthy_checks / enabled_checks * 100) if enabled_checks > 0 else 0,
            "memory_usage": latest_memory.value if latest_memory else None,
            "cpu_usage": latest_cpu.value if latest_cpu else None,
            "uptime_hours": (datetime.utcnow() - self.metrics.start_time).total_seconds() / 3600,
            "last_update": datetime.utcnow().isoformat()
        }